"""
    回测引擎主体

"""
import datetime

import pandas as pd

from .orders import Order
from .brokers import Broker
from .portfolio import PortTrack
from .dataLoader import DataLoader
from .paramsPipe import ParamsPipe
from .analysis import PortAnalysis

from data_resource.data_bases import engine


class BackTestEngine:
    """
    回测引擎主体
    方法：
        1. before_trading: 策略逻辑主体，需自行复写该类, 该方法需要传入当前数据帧参数
    params:
        data: 数据帧，columns=['trading', 'code', 'close', 'open', 'high', 'low', 'volume']
        cash: 初始资金, 默认为10000
    """

    def __init__(self, data: DataLoader, benchmark_data: pd.DataFrame = None, cash=10000):

        self.paramsPipe = ParamsPipe()
        self.data = data  # 数据管线
        self.benchmark = benchmark_data  # 基准数据
        self.current_row = 0
        self.total_portvalue = {}  # 记录策略净值曲线
        self.historical_picks = []  # 记录每期目标持仓

        self.port = PortTrack(data=self.data, params=self.paramsPipe, init_cash=cash)  # 创建组合类
        self.broker = Broker(data=data, port=self.port, cash=cash, params=self.paramsPipe)  # 创建券商类
        self.order = Order(data=data, broker=self.broker, port=self.port, params=self.paramsPipe)  # 创建订单类
        self.analysis = None  # 创建分析管线

        # ------- 通过管道传递参数 ------------ #
        self.today = None
        self.next_day = None
        self.paramsPipe.init_cash = cash  # 记录策略初始资金

    def __iter__(self):
        return self

    def __next__(self):
        # 逐帧数据执行策略逻辑
        if self.current_row >= len(self.data.unique_dates):
            raise StopIteration
        self.today = self.data.get_date(self.current_row)
        self.paramsPipe.update_today(self.today)
        self.next_day = self.data.get_date(self.current_row + 1)
        self.paramsPipe.update_nextDay(self.next_day)

        # 交易时为次日开盘价下单，故当日先进行组合净值更新
        self.total_portvalue.update({self.today: self.port.port_value})

        self.current_row += 1
        return self.current_row  # 迭代器返回当前数据帧索引

    def run(self):
        print("==================== 回测开始 ============================")
        start = datetime.datetime.now()
        for _ in range(len(self.data.unique_dates)):
            self.before_trading()
            self.__next__()
            self.after_trading()
            # 更新组合价值
            self.port.update_portfolio_value()
        end = datetime.datetime.now()
        print(f"========================= 回测结束，耗时{end - start} ===========================")

        if self.benchmark is not None:
            print("=========== 注入基准数据 ============")
            self.analysis = PortAnalysis(port_value=self.total_portvalue, benchmark=self.benchmark)
        else:
            self.analysis = PortAnalysis(port_value=self.total_portvalue)
        return

    def before_trading(self):
        # 策略逻辑主体，需自行复写该类; 执行时机为，每数据帧运行前
        pass

    def after_trading(self):
        # 策略逻辑主体，需自行复写该类; 执行时机为，每数据帧运行后
        pass

    def set_commision_fee(self, commision_fee: float):
        # 设置双边佣金费率，默认为3.5%。
        self.broker.set_commision_fee(commision_fee)

    def set_noh(self, noh: float):
        # 设置一手对应100股
        self.broker.set_noh(noh)

    def set_volume_limit(self, volume_limit: float):
        # 设置单笔成交量限额
        self.broker.set_volume_limit(volume_limit)

    def set_slip(self, slip: float):
        # 设置交易滑点，默认为0.1%
        self.broker.set_slip(slip)

    def get_pickups(self):
        """将历史选股结果输出"""
        _sql1 = """
            select ticker, short_name from quant_research.basic_info_stock;
        """
        _stocknames = pd.read_sql(_sql1, engine)

        # 获取历史记录
        results = []
        for x in self.historical_picks:
            _date = next(iter(x.keys()))
            _data = x[_date]
            _data['trade_date'] = _date
            results.append(_data)
        results = pd.concat(results)
        results.reset_index(drop=True, inplace=True)

        # 左连接股票名称
        results = pd.merge(results, _stocknames, left_on='code', right_on='ticker', how='left')
        results.drop(columns=['ticker'], inplace=True)
        return results
